#include "gamemessage.h"
#include "utils.h"
#include "networkevents.h"
#include "mrsa.h"
#include "maes.h"
#include "mbase64.h"


static const char* CMD_SYNCTIME = "SYNC_TIME";
static const char* CMD_PUBLIC_KEY = "PUBLIC_KEY";
static const char* CMD_PUBLIC_KEY_ACK = "PUBLIC_KEY_ACK";
static const char* CMD_HEARTBEAT = "HEARTBEAT";
static const char* CMD_HEARTBEAT_ACK = "HEARTBEAT_ACK";
static const char* CMD_CONNECT_DEVICE = "CONNECT_DEVICE";
static const char* CMD_CONNECT_DEVICE_ACK = "CONNECT_DEVICE_ACK";

#define ON_STATUS(s)    if (gmsg->on_status_cb) gmsg->on_status_cb(s);
#define ON_MESSAGE(cmd,data,code)    if (gmsg->on_message_cb) gmsg->on_message_cb(cmd,data,code);


enum MESSAGE_STATES {
    GMSGSTAT_NONE = 0, GMSGSTAT_IDLE, GMSGSTAT_WAIT_SYNCTIME, GMSGSTAT_HANDSHAKE,
    GMSGSTAT_CONNECTED, GMSGSTAT_DISCONNECT_WAIT, GMSGSTAT_CONNECTED_NOTIFY,
    GMSGSTAT_DEVICE_CONNECTING,
};

struct game_message {
    int fd;
    char serverip[128];
    int serverport;

    char version[32];

    int state;
    int connect_count;

    int64_t check_time;
    int64_t last_recvtime;
    int64_t heartbeat_time;

    uint8_t tmpbuf[1500];
    char command[64];
    char data[1100];
    int code;
    int64_t timestamp;
    int64_t timestamp_delta;

    char* sendbuf;
    char* sendbuf_tmp;
    int sendbuf_size;

    aes_context* aes;
    aes_context* aes_enc;
    uint8_t aes_key[24];
    rsa_context* rsa;
    char rsa_key[512];

    message_received_cb on_message_cb;
    message_status_cb on_status_cb;
};

static const char* AES_IV = "4b2d6c1a5e3h7g8f";
struct game_message* gmsg = NULL;

void gamemsg_process_command();

int my_rand(void* p)
{
    static int srand_inited = 0;
    if (!srand_inited) {
        srand_inited = 1;
        srand((uint32_t)time(NULL));
    }
    return rand();
}

int parse_message_packet(const char* input)
{
    const char *p1, *p2, *p3, *p4, *p;
    int n;

    p1 = strstr(input, "\"type\":");
    p2 = strstr(input, "\"data\":");
    p3 = strstr(input, "\"timestamp\":");
    p4 = strstr(input, "\"code\":");
    if (!p1 || !p2 || !p3) {
        LOGE("invalidate packet: %s", input);
        return 0;
    }

    p1 += 7;
    p2 += 7;
    p3 += 12;
    if (*p1 != '"' || (*p2 != '"' && *p2 != 'n') || (*p3 < '0' || *p3 > '9')) {
        LOGE("invalidate packet: %s", input);
        return 0;
    }

    p1++;
    p = p1;
    while (*p && *p != '"') p++;
    if (*p != '"') {
        LOGE("invalidate packet: %s", input);
        return -1;
    }
    n = p - p1;
    if (n > 60)
        n = 60;
    memcpy(gmsg->command, p1, n);
    gmsg->command[n] = 0;

    if (!strncmp(p2, "null", 4)) {
        gmsg->data[0] = 0;
    }
    else if (*p2 == '{') {
        p = p2 + 1;
        while (*p && *p != '}') p++;
        if (*p != '}') {
            LOGE("invalidate packet: %s", input);
            return -1;
        }
        n = p - p2;
        if (n > 1024)
            n = 1024;
        memcpy(gmsg->data, p2, n);
        gmsg->data[n] = 0;
    }
    else {
        p2++;
        p = p2;
        while (*p && *p != '"') p++;
        if (*p != '"') {
            LOGE("invalidate packet: %s", input);
            return -1;
        }
        n = p - p2;
        if (n > 1024)
            n = 1024;
        memcpy(gmsg->data, p2, n);
        gmsg->data[n] = 0;
    }

    gmsg->timestamp = atoll(p3);

    if (p4) {
        p4 += 7;
        gmsg->code = atoi(p4);
    }

    return 1;
}

int gamemsg_data_handler(const char* data, int size, void* param)
{
    int ret;
    if (size > 4) {
        uint8_t* p = (uint8_t*)data;
        int len = (((int)p[0]) << 24) | (((int)p[1]) << 16) | (((int)p[2]) << 8) | (int)p[3];

        if (len + 4 <= size) {
            if (len < 1200) {
                const char* input = data + 4;
                if (gmsg->aes) {
                    // AES解密
                    aes_setkey_dec(gmsg->aes, gmsg->aes_key, 128);
                    aes_crypt_cbc(gmsg->aes, AES_DECRYPT, len, (uint8_t*)AES_IV, (uint8_t*)input, (uint8_t*)gmsg->tmpbuf);
                    // 去掉AES加密时补的字节(ASCII: 1-16)
                    input = (const char*)gmsg->tmpbuf + (len - 1);
                    while (input > (const char*)gmsg->tmpbuf && *input != '}')
                        input--;
                    if (input == (const char*)gmsg->tmpbuf) {
                        LOGE("MESSAGE DECRYPT ERROR, len %d", len);
                        return len + 4;
                    }
                    input = (const char*)gmsg->tmpbuf;
                }

                ret = parse_message_packet(input);
                if (ret < 0)
                    return ret;
                if (ret > 0) {
                    gamemsg_process_command();
                }
            }
            return len + 4;
        }
    }
    return 0;
}

int gamemsg_connect_handler(int newfd, int errcode, void* param)
{
    if (errcode) {
        LOGE("connect server failed. error code %d", errcode);
        gmsg->state = GMSGSTAT_IDLE;
        gmsg->check_time = now_time2 + 2000;
        gmsg->last_recvtime = now_time2;
        gmsg->heartbeat_time = now_time2;
    }
    else {
        gmsg->fd = newfd;
        gmsg->state = GMSGSTAT_WAIT_SYNCTIME;
        gmsg->check_time = now_time2 + 1500;
        LOGI("Server connected. (%d)", newfd);
    }

    return 0;
}

int gamemsg_closed_handler(int errcode, void* param)
{
    LOGE("Server connection closed. error code %d", errcode);

    gamemsg_disconnect(0);

    if (gmsg->state == GMSGSTAT_HANDSHAKE)
        gmsg->check_time = now_time2 + 2000;
    else
        gmsg->check_time = now_time2 + 10000;
    gmsg->state = GMSGSTAT_IDLE;

    return 0;
}

int gamemsg_init()
{
    gmsg = (game_message*)malloc(sizeof(struct game_message));
    memset(gmsg, 0, sizeof(struct game_message));
    gmsg->fd = -1;
    gmsg->sendbuf_size = 1024 * 6;
    gmsg->sendbuf = (char*)malloc(gmsg->sendbuf_size);
    gmsg->sendbuf_tmp = (char*)malloc(gmsg->sendbuf_size);

    return 0;
}

void gamemsg_uninit()
{
    if (gmsg) {
        gamemsg_disconnect(0);
        if (gmsg->sendbuf)
            free(gmsg->sendbuf);
        if (gmsg->sendbuf_tmp)
            free(gmsg->sendbuf_tmp);
        free(gmsg);
    }
}

int gamemsg_connect(const char* serverip, int serverport, const char* version, message_received_cb message_cb, message_status_cb status_cb)
{
    if (gmsg->state >= GMSGSTAT_CONNECTED) {
        if (strcmp(serverip, gmsg->serverip) == 0 && serverport == gmsg->serverport) {
            gmsg->state = GMSGSTAT_CONNECTED_NOTIFY;
            return 0;
        }
    }
    gamemsg_disconnect(0);

    gmsg->fd = -1;
    strcpy(gmsg->serverip, serverip);
    gmsg->serverport = serverport;
    gmsg->on_message_cb = message_cb;
    gmsg->on_status_cb = status_cb;
    gmsg->state = GMSGSTAT_IDLE;
    gmsg->check_time = 0;
    strcpy(gmsg->version, version);

    return 0;
}

void gamemsg_disconnect(int delay_secs)
{
    if (!gmsg || gmsg->fd == -1)
        return;

    if (delay_secs > 0) {
        gmsg->state = GMSGSTAT_DISCONNECT_WAIT;
        gmsg->check_time = now_time2 + (delay_secs * 1000);
        return;
    }

    network_close(gmsg->fd);

    if (gmsg->aes) {
        free(gmsg->aes);
        gmsg->aes = NULL;

        free(gmsg->aes_enc);
        gmsg->aes_enc = NULL;
    }
    if (gmsg->rsa) {
        rsa_free(gmsg->rsa);
        free(gmsg->rsa);
        gmsg->rsa = NULL;
    }

    //memset(gmsg, 0, sizeof(struct game_message));
    gmsg->fd = -1;
}

void gamemsg_process()
{
    static int show_state = 0;
    int ret;

    if (!gmsg || gmsg->serverip[0] == 0)
        return;

    switch (gmsg->state) {
        case GMSGSTAT_IDLE:
            if (now_time2 < gmsg->check_time)
                return;
            show_state = 0;
            LOGI("connecting to %s:%d", gmsg->serverip, gmsg->serverport);
            ret = network_connect(gmsg->serverip, gmsg->serverport, NULL, 3000,
                                  gamemsg_connect_handler, gamemsg_data_handler, gamemsg_closed_handler, gmsg);
            if (ret < 0) {
                LOGE("connect game server %s:%d failed.(%d)", gmsg->serverip, gmsg->serverport, ret);
                gmsg->check_time = now_time2 + 3000;
                return;
            }

            gmsg->check_time = now_time2 + 5000;
            break;
        case GMSGSTAT_WAIT_SYNCTIME:
            if (now_time2 > gmsg->check_time) {
                gamemsg_disconnect(0);
                gmsg->check_time = now_time2 + 1000;
                gmsg->state = GMSGSTAT_IDLE;
                LOGE("can't receive SYNCTIME, try again later.");
            }
            break;
        case GMSGSTAT_HANDSHAKE:
            if (now_time2 > gmsg->check_time) {
                gamemsg_disconnect(0);
                LOGE("HANDSHAKE timeout.");
                gmsg->check_time = now_time2 + 1000;
                gmsg->state = GMSGSTAT_IDLE;
            }
            break;
        case GMSGSTAT_CONNECTED:
            // 检查发送心跳
            if (now_time2 - gmsg->last_recvtime > 10000) {
                gamemsg_disconnect(0);
                gmsg->check_time = 0;
                gmsg->state = GMSGSTAT_IDLE;
                ON_STATUS("RECV_TIMEOUT");
            }
            else if (now_time2 - gmsg->heartbeat_time >= 1000) {
                gamemsg_send(CMD_HEARTBEAT, "", 0);
                gmsg->heartbeat_time = now_time2;
            }
            break;
        case GMSGSTAT_DISCONNECT_WAIT:
            if (now_time2 > gmsg->check_time) {
                gamemsg_disconnect(0);
                LOGI("SERVER DISCONNECT!");
                gmsg->state = GMSGSTAT_NONE;	// 不做任何事
            }
            else if (now_time2 - gmsg->heartbeat_time >= 1000) {
                gamemsg_send(CMD_HEARTBEAT, "", 0);
                gmsg->heartbeat_time = now_time2;
            }
            break;
        case GMSGSTAT_CONNECTED_NOTIFY:
            gmsg->state = GMSGSTAT_CONNECTED;
            gmsg->last_recvtime = now_time2;
            ON_STATUS("CONNECTED");
            break;
        case GMSGSTAT_DEVICE_CONNECTING:
            if (now_time2 >= gmsg->check_time) {
                gamemsg_disconnect(0);
                gmsg->check_time = 0;
                gmsg->state = GMSGSTAT_IDLE;
                LOGE("CONNECT DEVICE timeout!");
            }
            break;
    }
}

void gamemsg_process_command()
{
    int len;

    if (strcmp(gmsg->command, CMD_HEARTBEAT_ACK) && strcmp(gmsg->command, CMD_PUBLIC_KEY_ACK))
        LOGI("RECV_MSG >> %s - %s (code:%d)", gmsg->command, gmsg->data, gmsg->code);

    gmsg->last_recvtime = now_time2;
    gmsg->heartbeat_time = now_time2;

    switch (gmsg->state) {
        case GMSGSTAT_WAIT_SYNCTIME:
            if (strcmp(gmsg->command, CMD_SYNCTIME) == 0) {
                gmsg->timestamp_delta = gmsg->timestamp - sys_gettime();
                gmsg->state = GMSGSTAT_HANDSHAKE;
                gmsg->check_time = now_time2 + 1500;

                if (!gmsg->rsa) {
                    // 生成RSA密钥
                    gmsg->rsa = (rsa_context*)malloc(sizeof(rsa_context));
                    rsa_init(gmsg->rsa, RSA_PKCS_V15, 0, my_rand, NULL);
                    rsa_gen_key(gmsg->rsa, 512, 65537);

                    // RSA公钥转换为十进制字符串
                    len = sizeof(gmsg->rsa_key);
                    mpi_write_string(&(gmsg->rsa->N), 10, gmsg->rsa_key, &len);
                    gmsg->rsa_key[len] = 0;
                }
                // 发送PUBLIC_KEY消息
                gamemsg_send(CMD_PUBLIC_KEY, gmsg->rsa_key, 0);
            }
            else {
                gamemsg_disconnect(0);
                gmsg->check_time = now_time2 + 2000;
                ON_STATUS("NOT_SYNCTIME");
            }
            break;
        case GMSGSTAT_HANDSHAKE:
            if (strcmp(gmsg->command, CMD_PUBLIC_KEY_ACK) == 0) {
                len = 256;
                if (base64_decode(gmsg->tmpbuf, &len, (uint8_t*)gmsg->data, strlen(gmsg->data)) == 0) {
                    len = sizeof(gmsg->data);
                    if (rsa_pkcs1_decrypt(gmsg->rsa, RSA_PRIVATE, &len, gmsg->tmpbuf, (uint8_t*)gmsg->data) == 0) {
                        if (len > 16)
                            len = 16;
                        memcpy(gmsg->aes_key, gmsg->data, len);
                        gmsg->data[len] = 0;
                        gmsg->aes = (aes_context *)malloc(sizeof(aes_context));
                        gmsg->aes_enc = (aes_context *)malloc(sizeof(aes_context));
                        gmsg->state = GMSGSTAT_CONNECTED;
                        gmsg->check_time = 0;
                        ON_STATUS("CONNECTED");
                    }
                    else {
                        LOGE("HANDSHAKE DECRYPT ERROR");
                        gamemsg_disconnect(0);
                        gmsg->check_time = now_time2 + 2000;
                        ON_STATUS("PUBLIC_KEY_ACK_DECRYPT_ERROR");
                    }
                }
                else {
                    LOGE("HANDSHAKE ERROR: %s", gmsg->data);
                    gamemsg_disconnect(0);
                    gmsg->check_time = now_time2 + 2000;
                    ON_STATUS("PUBLIC_KEY_ACK_BASE64_ERROR");
                }
            }
            else {
                gamemsg_disconnect(0);
                gmsg->check_time = now_time2 + 2000;
                ON_STATUS("NOT_PUBLIC_KEY_ACK");
            }

            rsa_free(gmsg->rsa);
            free(gmsg->rsa);
            gmsg->rsa = NULL;
            break;
        case GMSGSTAT_CONNECTED:
        case GMSGSTAT_DEVICE_CONNECTING:
            if (strcmp(gmsg->command, CMD_HEARTBEAT_ACK)) {
                ON_MESSAGE(gmsg->command, gmsg->data, gmsg->code);
                if (gmsg->state == GMSGSTAT_DEVICE_CONNECTING
                        && strcmp(gmsg->command, CMD_CONNECT_DEVICE_ACK) == 0) {
                    gmsg->state = GMSGSTAT_CONNECTED;
                }
            }
            break;
    }
}

void gamemsg_send(const char* command, const char* data, int code)
{
    if (gmsg->fd != -1) {
        int64_t timestamp = gmsg->timestamp_delta + sys_gettime();
        int len;
        char* sendbuf;

        if (!data)
            data = "";

        len = strlen(data);
        if (len + 200 >= gmsg->sendbuf_size) {
            char* tmp = (char*)malloc(len + 200);
            char* tmp1 = (char*)malloc(len + 200);

            if (!tmp || !tmp1) {
                if (tmp)
                    free(tmp);
                LOGE("SEND_MESSAGE ERROR: failed to alloc memory!");
                return;
            }

            free(gmsg->sendbuf);
            free(gmsg->sendbuf_tmp);
            gmsg->sendbuf = tmp;
            gmsg->sendbuf_tmp = tmp1;
            gmsg->sendbuf_size = len + 200;
        }
        sendbuf = gmsg->sendbuf + 4;

        if (gmsg->aes_enc) {
            // AES加密消息
            char *tmpbuf = gmsg->sendbuf_tmp;
            int n;

            if (data && *data == '{') {
                n = sprintf(tmpbuf,
                            "{\"type\":\"%s\",\"timestamp\":%lld,\"data\":%s,\"code\":%d,\"version\":\"%s\"}",
                            command, timestamp, data, code,  gmsg->version);
            }
            else {
                n = sprintf(tmpbuf,
                            "{\"type\":\"%s\",\"timestamp\":%lld,\"data\":\"%s\",\"code\":%d,\"version\":\"%s\"}",
                            command, timestamp, data, code,  gmsg->version);
            }

            if (strcmp(command, CMD_HEARTBEAT) && strcmp(command, CMD_PUBLIC_KEY))
                LOGI("SEND >> %s", tmpbuf);

            // AES加密数据长度为16的倍数
            // PKCS5Padding：填充的原则是，如果长度少于16个字节，需要补满16个字节，
            //               如果字符串长度正好是16字节，则需要再补16个字节的十进制的16。
            if ((n & 0x0f) > 0)
                len = ((n / 16) + 1) * 16;
            else
                len = n + 16;
            // 加密时补的是十进制1到16
            while (n < len)
                tmpbuf[n++] = 0x01;

            // AES加密
            aes_setkey_enc(gmsg->aes_enc, gmsg->aes_key, 128);
            aes_crypt_cbc(gmsg->aes_enc, AES_ENCRYPT, len, (uint8_t*)AES_IV, (uint8_t*)tmpbuf, (uint8_t*)sendbuf);
        }
        else {
            // 不加密消息
            len = sprintf(sendbuf,
                          "{\"type\":\"%s\",\"timestamp\":%lld,\"data\":\"%s\",\"code\":%d,\"version\":\"%s\"}",
                          command, timestamp, data, code, gmsg->version);

            if (strcmp(command, CMD_HEARTBEAT) && strcmp(command, CMD_PUBLIC_KEY))
                LOGI("SEND >> %s", sendbuf);
        }

        sendbuf -= 4;
        sendbuf[0] = (uint8_t)((len >> 24));
        sendbuf[1] = (uint8_t)((len >> 16) & 0xff);
        sendbuf[2] = (uint8_t)((len >> 8) & 0xff);
        sendbuf[3] = (uint8_t)(len & 0xff);
        network_send(gmsg->fd, sendbuf, len + 4);

        if (strcmp(command, CMD_CONNECT_DEVICE) == 0) {
            gmsg->state = GMSGSTAT_DEVICE_CONNECTING;
            gmsg->check_time = now_time2 + 5000;
        }
    }
}

int64_t gamemsg_get_timestamp_delta()
{
    if (gmsg)
        return gmsg->timestamp_delta;
    return 0;
}
